import numpy as np
import os
import scipy.signal as signal
import matplotlib.pyplot as plt
from scipy.signal import butter, filtfilt
from sklearn.preprocessing import StandardScaler

# 非线性特征
def hilbert(data):
    # 对每个通道的信号进行希尔伯特变换
    analytic_signals = signal.hilbert(data, axis=1)  # 沿着第二个轴 (时间轴) 进行变换
    # 提取特征
    instantaneous_phase = np.angle(analytic_signals)  # 瞬时相位
    instantaneous_amplitude = np.abs(analytic_signals)  # 瞬时幅值
    return instantaneous_amplitude


# 频域特征
# 短时傅里叶变换
def extract_quick_FFT(signals, fs=400, window_length=128, hop_length=64):
    '''
    计算短时傅里叶变换
    选窗：'hann', 'hamming', 'blackman', 窗函数能够更好地抑制边界效应并提高结果的可靠性
    '''
    f, t, Zxx = signal.stft(
        signals, 
        fs=fs, 
        window='hamming', 
        nperseg=window_length, 
        noverlap=window_length - hop_length
    )
    #  可视化时频图
    # Zxx_channel = np.abs(Zxx[10, :, :]).T
    # plt.pcolormesh(Zxx_channel, shading='gouraud')  # 注意这里需要使用 Zxx_channel.T 进行转置
    # plt.xlabel('Time (s)')
    # plt.ylabel('Frequency (Hz)')
    # plt.title('STFT')
    # plt.xlim(0, 15)
    # plt.savefig(f'sfft_{band}.png')
    nergy_mean = np.mean(np.abs(Zxx)**2, axis=len(signals.shape)-1)
    return nergy_mean

# 功率谱密度 (PSD): 反映信号在不同频率上的功率分布
def extract_psd_features(signals, fs):
    #  使用 Welch 方法估计功率谱密度
    f, psd = signal.welch(signals, fs=fs)
    return f, psd


# 滤波 分别对不同的波段进行预处理
def butter_bandpass(lowcut, highcut, fs, order=5):
    nyq = 0.5 * fs
    low = lowcut / nyq
    high = min(highcut / nyq, 0.999)
    b, a = butter(order, [low, high], btype='bandpass')
    return b, a  # 返回滤波器的系数

def butter_bandpass_filter(data, lowcut, highcut, fs, order=5):
    b, a = butter_bandpass(lowcut, highcut, fs, order=order)
    y = filtfilt(b, a, data)
    return y

def extract_part(data, sr, windowLength, class_start, dir_path, file):
    class_name = class_start
    for win in range(windowLength):
        start = int(np.floor(win * sr))
        stop = int(np.floor(start + sr))
        feat = data[:, start:stop]
        np.save(os.path.join(dir_path, file[:-4]+'_'+str(class_name)), feat)
        class_name += 1

# segmentation
def segment_pooling(arr):
    _, rate = arr.shape
    winL = 0.1  # 每个片段 100ms
    overlap = 0.05
    length = int(15 / (winL-overlap)) - 1    # 将 15s 的信号切割成 length 段
    segment = int(rate//length)   # 每段有40个点 segment * 2
    
    pooled_subarrays = []
    for j in range(length):
        xi = arr[:, j*segment : (j+2)*segment]
        pi = np.mean(xi, axis=1)
        pooled_subarrays.append(pi)
    concatenated = np.stack(pooled_subarrays, axis=1)
    result = np.array(concatenated)
    return result


if __name__ == '__main__':
    # 划窗，窗口长度为1s
    winL = 15
    duration = 15
    fs = 400
    cls_num = 2
    root_data_path = '/root/data/video_decoding'
    eeg_path = os.path.join(root_data_path, 'cut_data_zscore')      # original file path
    outputs = os.path.join(root_data_path, f'npy_data_class_{cls_num}')   # output feature
    time_files = os.listdir(eeg_path)
    transfer = StandardScaler()
    
    # 定义波段
    max_band = min(fs, 200)
    bands = {
            'theta': (4, 7),
            'alpha': (8, 12),
            'beta': (13, 30),
            'low_gamma': (31, 90),
            'high_gamma': (90, max_band)
        }
    for data_id in ['train', 'val']:
        file_names = os.listdir(os.path.join(eeg_path, data_id))
        for file in file_names:
            # neural data
            seeg_data = np.load(os.path.join(eeg_path, data_id, file))
            print(file, ' is load')

            for band in bands:
                low, high = bands[band]
                filter_feats = butter_bandpass_filter(seeg_data, low, high, fs, order=5)
                
                # 绘制前 5 个信道的数据，并将其绘制在同一张图上
                # plt.figure(figsize=(12, 6)) # 设置画布大小
                # for i in range(5):  
                #     plt.subplot(5, 1, i+1)  # 创建子图，5 行 1 列，当前子图为第 i+1 个
                #     plt.plot(filter_feats[i, :]) # 绘制第 i 个信道的信号
                #     plt.title(f'Channel {i+1}')  # 设置子图标题
                #     plt.xlabel('Sample Number') # 设置横轴标签
                #     plt.ylabel('Signal Amplitude') # 设置纵轴标签
                # plt.tight_layout() # 调整子图间距
                # plt.savefig(f'{band}.png')
                
                # 去除线性趋势
                if cls_num == 2:
                    filter_feats = signal.detrend(seeg_data, axis=1, type='constant')   # linear 更差
                else:
                    filter_feats = signal.detrend(seeg_data, axis=1)
                    
                # normalization
                feature = transfer.fit_transform(filter_feats)
                
                store_path = os.path.join(outputs, band)
                if not os.path.exists(store_path):
                    os.mkdir(store_path)
                store_path = os.path.join(store_path, data_id)
                if not os.path.exists(store_path):
                    os.mkdir(store_path)
                    
                if cls_num == 2:
                    # feature = extract_quick_FFT(feature, fs, window_length=128, hop_length=64)    # 性能下降
                    # feature = segment_pooling(feature)        除了 low_gamma 其他反而下降了
                    store_path = os.path.join(store_path, file)
                    np.save(store_path, feature)
                else:
                    feature = hilbert(feature)
                    if file[16]=='1':
                        class_start = 0
                    elif file[16]=='2':
                        class_start = 15
                    extract_part(filter_feats, fs, winL, class_start, store_path, file)            
                